import os
import cv2
import numpy as np
import tensorflow as tf
import label

# 设置训练集、测试集的地址
base_dir = './'
train_dir = os.path.join(base_dir, 'Training')
test_dir = os.path.join(base_dir, 'test')
# test_multiple_dir = os.path.join(base_dir, 'test-multiple_fruits')
saved_files = './output_files'

# 创建模型结果文件
if not os.path.exists(saved_files):
    os.makedirs(saved_files)

# 获取训练集目录及种类数
# labels = os.listdir(train_dir)
labels = label.getFruitLabel()
num_classes = len(labels)


def test_model(name=""):
    model_out_dir = os.path.join(saved_files, f'{name}.h5')

    if not os.path.exists(model_out_dir):
        print("No saved model found")
        exit(0)
    # 读取模型
    model = tf.keras.models.load_model(model_out_dir)

    # 获取图片并处理
    # image = cv2.imread(test_dir + '/Orange/Orange003011.png')
    image = cv2.imread('./0035.jpg')
    image = cv2.resize(image, (100, 100))

    # 设置图片维度
    data = np.ndarray(shape=(1, 100, 100, 3), dtype=np.int_)
    image_array = np.asarray(image)
    data[0] = image_array

    # 预测结果
    result = model.predict(data, 1)
    print("预测结果为: " + labels[result.argmax(axis=-1)[0]])


test_model(name='test')
